import * as React from 'react'

export interface ModalsContextValue<
  TModals extends Record<string, React.ComponentType<any>> = Record<
    string,
    React.ComponentType<any>
  >,
  TTypes extends Extract<keyof TModals, string> = Extract<
    keyof TModals,
    string
  >,
> {
  open: <T extends OpenOptions<TTypes>>(
    componentOrOptions: T extends {
      component: infer TComponent extends React.ComponentType<any>
    }
      ? WithModalOptions<React.ComponentPropsWithRef<TComponent>>
      : T extends {
            type: infer TType extends keyof TModals
          }
        ? WithModalOptions<React.ComponentPropsWithRef<TModals[TType]>>
        : T,
    options?: T extends React.FC<any>
      ? WithModalOptions<React.ComponentPropsWithRef<T>>
      : never,
  ) => ModalId
  alert: <T extends AlertDialogOptions>(options: T) => ModalId
  confirm: <T extends ConfirmDialogOptions>(options: T) => ModalId
  close: (id: ModalId) => void
  closeAll: () => void
}

export const ModalsContext = React.createContext<ModalsContextValue<
  Record<string, React.ComponentType<any>>
> | null>(null)

export interface ModalsProviderProps<
  TModals extends Record<string, React.ComponentType<any>> = Record<
    string,
    React.FC<any>
  >,
> {
  children: React.ReactNode
  modals?: TModals
  render?: (props: ModalsContextValue<TModals>) => React.ReactNode
}

export type ModalId = string | number

type WithModalOptions<T> = Omit<T, 'open' | 'onOpenChange'> & ModalOptions

export interface ModalOptions {
  title?: string
  body?: React.ReactNode
  open?: boolean
  onOpenChange?: (details: { open: boolean }) => void
  onClose?: (args: { force?: boolean }) => Promise<boolean | undefined> | void
  [key: string]: any
}

export interface AlertDialogOptions extends ModalOptions {
  onConfirm?: () => Promise<void> | void
}

export interface ConfirmDialogOptions extends ModalOptions {
  leastDestructiveFocus?: 'cancel' | 'confirm'
  onConfirm?: () => Promise<void> | void
  onCancel?: () => Promise<void> | void
}

export interface OpenOptions<TModalTypes extends string> extends ModalOptions {
  type?: TModalTypes
  scope?: ModalScopes
}

export type ModalScopes = 'modal' | 'alert'

export interface ModalConfig<
  TModalOptions extends ModalOptions = ModalOptions,
  TModalTypes extends string = string,
> {
  /**
   * The modal id, autogenerated when not set.
   * Can be used to close modals.
   */
  id?: ModalId | null
  /**
   * The modal props
   */
  props?: TModalOptions | null
  /**
   * The modal scope
   * Modals can only have one level per scope.
   * The default scopes are 'modal' and 'alert', alerts can be openend above modals.
   */
  scope?: ModalScopes | string
  /**
   * The modal type to open.
   * Build in types are 'modal', 'drawer', 'alert', 'confirm'
   *
   * Custom types can be configured using the `modals` prop of `ModalProvider`
   */
  type?: TModalTypes
  /**
   * Render a custom modal component.
   * This will ignore the `type` param.
   */
  component?: React.ComponentType<ModalOptions>
  /**
   * Whether the modal is open or not.
   * This is used internally to keep track of the modal state.
   */
  open?: boolean
}

const initialModalState: ModalConfig = {
  id: null,
  props: null,
  type: 'modal',
}

export function ModalsProvider({ children, modals }: ModalsProviderProps) {
  // Note that updating the Set doesn't trigger a re-render,
  // use in conjuction with setActiveModals
  const _instances = React.useMemo(() => new Set<ModalConfig>(), [])

  const [activeModals, setActiveModals] = React.useState<
    Record<string, ModalConfig>
  >({
    modal: initialModalState,
  })

  const getModalComponent = React.useMemo(() => {
    const _modals: Record<string, React.ComponentType<any>> = modals || {}

    return (type = 'modal') => {
      const component = _modals[type] || _modals.modal

      return component
    }
  }, [modals])

  const setActiveModal = React.useCallback(
    (modal: ModalConfig, scope?: string) => {
      if (!scope) {
        return setActiveModals({
          modal,
        })
      }
      setActiveModals((prevState) => ({
        ...prevState,
        [scope]: modal,
      }))
    },
    [],
  )

  const open = React.useCallback(
    <T extends OpenOptions<any>>(
      componentOrOptions: any,
      options?: T extends React.FC<any>
        ? WithModalOptions<React.ComponentPropsWithRef<T>>
        : never,
    ): ModalId => {
      let _options: ModalOptions
      if (typeof componentOrOptions === 'function') {
        _options = {
          component: componentOrOptions,
          ...options,
        } as unknown as T
      } else {
        _options = componentOrOptions
      }

      const {
        id = _instances.size + 1,
        type = 'modal',
        scope = 'modal',
        component,
        ...props
      } = _options

      const modal: ModalConfig<T> = {
        id,
        props: props as T,
        type,
        scope,
        component,
        open: true,
      }

      _instances.add(modal)
      setActiveModal(modal, scope)

      return id
    },
    [_instances],
  )

  const alert = React.useCallback(
    (options: ConfirmDialogOptions) => {
      return open({
        ...options,
        type: 'alert',
        scope: 'alert',
      })
    },
    [open],
  )

  const confirm = React.useCallback(
    (options: ConfirmDialogOptions) => {
      return open({
        ...options,
        type: 'confirm',
        scope: 'alert',
      })
    },
    [open],
  )

  const closeComplete = React.useCallback((id?: ModalId | null) => {
    const modals = [...Array.from(_instances)]
    const modal = modals.filter((modal) => modal.id === id)[0]

    _instances.delete(modal)

    const scoped = modal && modals.filter(({ scope }) => scope === modal.scope)

    if (scoped?.length === 1) {
      setActiveModal(initialModalState, modal.scope)
    }
  }, [])

  const close = React.useCallback(
    async (id?: ModalId | null, force?: boolean) => {
      const modals = [...Array.from(_instances)]
      const modal = modals.filter((modal) => modal.id === id)[0]

      if (!modal) {
        return
      }

      const shouldClose = await modal.props?.onClose?.({ force })
      if (shouldClose === false) {
        return
      }

      const scoped = modals.filter(({ scope }) => scope === modal.scope)

      if (scoped.length === 1) {
        setActiveModal(
          {
            ...modal,
            open: false,
          },
          modal.scope,
        )
      } else if (scoped.length > 1) {
        setActiveModal(scoped[scoped.length - 2], modal.scope)
      } else {
        setActiveModal(
          {
            id: null,
            props: null,
            type: modal.type, // Keep type same as last modal type to make sure the animation isn't interrupted
          },
          modal.scope,
        )
      }

      // @todo this is not ideal, but not all modals support onCloseComplete
      setTimeout(() => closeComplete(id), 200)
    },
    [closeComplete],
  )

  const closeAll = React.useCallback(() => {
    _instances.forEach((modal) => modal.props?.onClose?.({ force: true }))
    _instances.clear()

    setActiveModal(initialModalState)
  }, [_instances])

  const context = React.useMemo(
    () => ({
      open,
      alert,
      confirm,
      close,
      closeAll,
    }),
    [open, alert, confirm, close, closeAll],
  )

  const content = React.useMemo(
    () =>
      Object.entries(activeModals).map(([scope, config]) => {
        const Component = config.component || getModalComponent(config.type)

        const { title, body, ...props } = config.props || {}

        return (
          <Component
            key={scope}
            title={title}
            {...props}
            open={!!config.open}
            onOpenChange={(details) =>
              details.open === false && close(config.id)
            }
            onExitComplete={() => closeComplete(config.id)}
          >
            {body}
          </Component>
        )
      }),
    [activeModals, getModalComponent],
  )

  return (
    <ModalsContext.Provider value={context}>
      {content}
      {children}
    </ModalsContext.Provider>
  )
}

export const useModalsContext = () => React.useContext(ModalsContext)

export const useModals = () => {
  const modals = useModalsContext()

  if (!modals) {
    throw new Error('useModals must be used within a ModalsProvider')
  }

  return modals
}
